# -*- coding: utf-8 -*-
"""
@Time ： 2020-11-19 14:24
@Auth ： lixin
@Description：时间序列分析训练类

"""
from statsmodels.tsa.arima_model import ARIMA
from statsmodels.tsa.stattools import adfuller
# from sympy.physics.quantum.tests.test_pauli import sm
import statsmodels.api as sm

from algo.Algo_interface import Algo_interface
# from api import arima_pre_recovery
# import sys
# sys.path.append("..")
# import api



class TimeSeriesAnalysis(Algo_interface):
    def __init__(self, model_type, model_name, model_params):
        self.task_type = model_type
        self.model_name = model_name
        self.model_params = model_params
        self.model = None
        self.build_model()
        # return self.model

    def set_model(self, model):
        self.model = model
        return 1

    def get_model(self):
        return self.model

    def build_model(self):
        pass
    def train(self,data):
        stander_predict=[]
        x_train=data
        match_p = False
        diff_data = []  # 用于差分还原
        diff = x_train
        if 'd' in self.model_params:  # 用户指定了d，则直接使用这个d作为差分次数
            d = self.model_params['d']
            diff_data.append(diff)
            match_p = True
            for i in range(d):
                diff = diff.diff(1)
                diff = diff.dropna()
        else:
            for d in range(3):
                dftest = adfuller(diff)
                diff_data.append(diff)
                if dftest[1] < 0.05:
                    match_p = True
                    break
                else:
                    diff = diff.diff(1)
                    diff = diff.dropna()

        if match_p:
            if 'd' in self.model_params:
                self.model_params.pop('d')
            self.model_params['ic'] = 'aic'
            result = sm.tsa.arma_order_select_ic(diff, **self.model_params)['aic_min_order']
            order = (result[0], d, result[1])
            self.model = ARIMA(x_train, order=order).fit()
            predict_ts = self.model.predict()
            # 差分还原
            if len(diff_data) > 1:
                diff_data = diff_data[:-1]
                stander_predict = self.arima_pre_recovery(predict_ts, diff_data, 1)
            else:
                stander_predict = predict_ts

        else:
            # 数据查分后依然无法平滑时写入异常
            print("数据查分后依然无法平滑时写入异常")
        return stander_predict
    def predict(self,time_step):
        y_pred = self.model.forecast(int(time_step))[0]  #返回3个array，第一个为预测结果
        return y_pred

    def arima_pre_recovery(pre, diff_data, step=1):
        """
        功能：还原arima预测结果
        :param pre: pd.Series, 预测值
        :param diff_data: list, 如[origin, t1] origin:原始数据, t1:第一次差分结果
        :param step: 差分步长, 默认为1
        :return: rec
        """
        rec = pre  # 初始化还原结果
        diff_data.reverse()
        for d in diff_data:  # 逐阶还原
            d_shift = d.shift(step)
            rec = rec.add(d_shift)
        rec = rec.dropna()
        return rec